# -*- coding: utf-8 -*-
import sys
import os
import copy
import re
import json
from typing import List
from collections import defaultdict


# 标准事件：完整下发，涵盖这个时间的时间范围内的其他
STANDARD_EVENTS = ["[AclCompileAndExecute]", "[OpExecute]"]
DEVICE_EVENT = "[Compute]"
DEV_HOST_CPU = "Host CPU"
DEV_DEVICE = "Device"


class Record:
    def __init__(self):
        self.node_name: str = None
        self.event: str = None
        self.et: str = None
        self.tid: int = None
        self.timestamp: int = None
        self.iteration: str = None
        self.device: str = DEV_HOST_CPU


def get_name(rec: Record):
    if rec.node_name is None:
        return rec.event
    else:
        return "{}@{}".format(rec.node_name, rec.event)


class EventRecode(Record):
    def __init__(self):
        super().__init__()
        self.start: int = None
        self.end: int = None

    def __str__(self):
        return "{}, Thread {}@{}, start {}, end {}".format(
            get_name(self), self.tid, self.device, self.start, self.end
        )

    @staticmethod
    def create_from_rec(rec: Record):
        er = EventRecode()
        er.__dict__ = copy.deepcopy(rec.__dict__)
        return er

class ProfilingData:
    def __init__(self):
        self.version: str = "1.0"
        self.records: List[Record] = []
        self.event_records: List[EventRecode] = []

    def add_record(self, rec: Record):
        self.records.append(rec)

    def add_event_record(self, rec: EventRecode):
        self.event_records.append(rec)


class ProfilingDataAnalyzer():
    V1_START_RE = re.compile(r'Profiler version: (?P<version>[\d+\.]+), dump start, records num: \d+')
    V1_RE = re.compile(
        r'(?P<timestamp>\d+) (?P<tid>\d+) (?P<element>(\[\S+\])|(UNKNOWN\(-?\d+\))) (?P<event>(\[\S+\])|(UNKNOWN\(-?\d+\))) (?P<et>(Start)|(End)|(UNKNOWN\(-?\d+\)))')

    def __init__(self, file_path):
        self._file_path = file_path

    def read_in_records(self):
        pds = []
        with open(self._file_path, 'r') as f:
            pd = ProfilingData()
            for line in f:
                ma = ProfilingDataAnalyzer.V1_START_RE.match(line)
                if ma is not None:
                    pd.version = ma.group('version')
                    continue
                if line.startswith("Profiling dump end"):
                    pds.append(pd)
                    pd = ProfilingData()
                    continue
                ma = ProfilingDataAnalyzer.V1_RE.match(line)
                if ma is not None:
                    rec = Record()
                    rec.timestamp = int(ma.group("timestamp"))
                    rec.tid = int(ma.group("tid"))
                    rec.node_name = ma.group("element")
                    if rec.node_name.startswith("UNKNOWN"):
                        rec.node_name = None
                    rec.event = ma.group("event")
                    if rec.event == DEVICE_EVENT:
                        rec.device = DEV_DEVICE
                        rec.tid = 1
                    rec.et = ma.group("et")
                    pd.add_record(rec)
                    continue
        return pds

    @staticmethod
    def read_in_event_records(pds: [ProfilingData]):
        for pd in pds:
            records = pd.records[:]
            records.sort(key=lambda rec: rec.timestamp)
            keys_to_starts = defaultdict(list)
            for rec in records:
                name = get_name(rec)
                if rec.et == 'Start':
                    keys_to_starts[name].append(rec)
                    continue
                if rec.et == 'End':
                    if len(keys_to_starts[name]) == 0:
                        print("WARNING: Drop record {}, because can not find a start record for it.".format(rec))
                        continue
                    start_rec = keys_to_starts[name].pop()
                    er = EventRecode.create_from_rec(start_rec)
                    er.start = start_rec.timestamp
                    er.end = rec.timestamp
                    pd.add_event_record(er)
            for starts in keys_to_starts.values():
                for no_end_start_rec in starts:
                    print("WARNING: Drop record {}, because can not find a end event for it.".format(no_end_start_rec))

    def read_in_profiling_file(self):
        pds = self.read_in_records()
        ProfilingDataAnalyzer.read_in_event_records(pds)
        return pds


class TracingReporter():
    def __init__(self, pds: [ProfilingData]):
        self._pds = pds

    def report(self, file_path):
        for i, pd in enumerate(self._pds):
            j_file = []
            for rec in pd.event_records:
                j_rec = {
                    "name": get_name(rec),
                    "ph": "X",
                    "pid": rec.device,
                    "tid": rec.tid,
                    "ts": rec.start / 1000,
                    "dur": (rec.end - rec.start) / 1000
                }
                j_file.append(j_rec)

            dump_path = "{}_tracing_{}{}".format(file_path, i, ".json")
            print("Dump to tracing file {}.".format(dump_path))
            with open(dump_path, "w") as f:
                json.dump(j_file, f)


def main(input_path, output_dir):
    analyzer = ProfilingDataAnalyzer(input_path)
    pds = analyzer.read_in_profiling_file()

    reporter = TracingReporter(pds)
    reporter.report(output_dir)


if __name__ == '__main__':
    input_file = sys.argv[1]
    input_name = os.path.basename(input_file)
    input_name = os.path.splitext(input_name)[0]
    main(sys.argv[1], './' + str(input_name))
